import requests
import cv2
import os
import re
from urllib.parse import unquote
from flask import Flask, request, jsonify, send_file
from flask_cors import CORS
import shutil
import uuid

app = Flask(__name__)
CORS(app)  # 启用CORS

# 配置文件存储路径
UPLOAD_FOLDER = 'temp'
FRAMES_FOLDER = 'frames'
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(FRAMES_FOLDER, exist_ok=True)

def download_video(url, save_path):
    """下载视频"""
    headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36',
        'Referer': 'https://www.douyin.com/',
        'Accept': '*/*',
        'Accept-Encoding': 'gzip, deflate, br',
        'Accept-Language': 'zh-CN,zh;q=0.9,en;q=0.8',
        'Connection': 'keep-alive'
    }
    try:
        print(f"开始下载视频: {url}")
        response = requests.get(url, headers=headers, stream=True, verify=False)
        
        if response.status_code == 200:
            with open(save_path, 'wb') as f:
                for chunk in response.iter_content(chunk_size=1024):
                    if chunk:
                        f.write(chunk)
            print("视频下载成功")
            return True
        else:
            print(f"下载失败，状态码: {response.status_code}")
            return False
    except Exception as e:
        print(f"下载出错: {str(e)}")
        return False

def get_video_url_from_share(share_url):
    """从分享链接获取视频实际地址"""
    try:
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36',
            'Referer': 'https://www.douyin.com/'
        }
        
        # 处理重定向获取真实URL
        response = requests.get(share_url, headers=headers, allow_redirects=True)
        real_url = response.url
        
        # 提取视频ID
        video_id = re.search(r'/video/(\d+)', real_url)
        if not video_id:
            print("无法从URL中提取视频ID")
            return None
            
        video_id = video_id.group(1)
        api_url = f'https://www.iesdouyin.com/web/api/v2/aweme/iteminfo/?item_ids={video_id}'
        
        # 获取视频信息
        response = requests.get(api_url, headers=headers)
        if response.status_code != 200:
            print(f"获取视频信息失败: {response.status_code}")
            return None
            
        data = response.json()
        if 'item_list' not in data or not data['item_list']:
            print("未找到视频信息")
            return None
            
        video_url = data['item_list'][0]['video']['play_addr']['url_list'][0]
        video_url = video_url.replace('playwm', 'play')  # 去水印
        
        return video_url
    except Exception as e:
        print(f"解析分享链接出错: {str(e)}")
        return None

def extract_frames(video_path, output_dir):
    """提取视频帧"""
    print(f"开始提取帧，视频路径: {video_path}")
    try:
        cap = cv2.VideoCapture(video_path)
        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        frames_saved = 0

        for i in range(frame_count):
            ret, frame = cap.read()
            if not ret:
                break
            frame_path = os.path.join(output_dir, f"frame_{i}.jpg")
            cv2.imwrite(frame_path, frame)
            frames_saved += 1

        cap.release()
        print(f"成功提取 {frames_saved} 帧")
        return frames_saved
    except Exception as e:
        print(f"提取帧出错: {str(e)}")
        return 0

@app.route('/extract-frames', methods=['GET'])
def extract_frames_api():
    try:
        # 获取视频URL参数
        video_url = request.args.get('video_url')
        if not video_url:
            return jsonify({'error': '请提供视频URL'}), 400

        print(f"处理视频URL: {video_url}")
        
        # 如果是分享链接，先获取实际视频地址
        if 'v.douyin.com' in video_url or 'www.douyin.com' in video_url:
            video_url = get_video_url_from_share(video_url)
            if not video_url:
                return jsonify({'error': '无法获取视频地址'}), 400
        
        # 创建唯一的会话ID和相关目录
        session_id = str(uuid.uuid4())
        session_folder = os.path.join(FRAMES_FOLDER, session_id)
        os.makedirs(session_folder, exist_ok=True)
        
        video_path = os.path.join(UPLOAD_FOLDER, f"{session_id}.mp4")

        # 下载视频
        if not download_video(video_url, video_path):
            return jsonify({'error': '视频下载失败'}), 400

        # 提取帧
        frames_count = extract_frames(video_path, session_folder)
        
        if frames_count > 0:
            # 创建压缩文件
            zip_path = os.path.join(UPLOAD_FOLDER, f"{session_id}.zip")
            shutil.make_archive(zip_path[:-4], 'zip', session_folder)
            
            # 清理临时文件
            os.remove(video_path)
            shutil.rmtree(session_folder)
            
            # 返回压缩文件
            return send_file(
                zip_path,
                mimetype='application/zip',
                as_attachment=True,
                download_name='frames.zip'
            )
        else:
            return jsonify({'error': '帧提取失败'}), 500

    except Exception as e:
        print(f"处理请求时出错: {str(e)}")
        return jsonify({'error': str(e)}), 500

@app.route('/health', methods=['GET'])
def health_check():
    return jsonify({'status': 'healthy'})

def main():
    # 创建保存目录
    if not os.path.exists(FRAMES_FOLDER):
        os.makedirs(FRAMES_FOLDER)
        
    # 获取视频链接
    video_url = input('请输入抖音视频链接：').strip()
    
    # 如果是分享链接，先获取实际视频地址
    if 'v.douyin.com' in video_url or 'www.douyin.com' in video_url:
        video_url = get_video_url_from_share(video_url)
        if not video_url:
            print("无法获取视频地址")
            return
    
    # 下载视频
    video_path = 'temp_video.mp4'
    if not download_video(video_url, video_path):
        print("视频下载失败")
        return
    
    # 提取帧
    frames_count = extract_frames(video_path, FRAMES_FOLDER)
    
    # 清理临时文件
    if os.path.exists(video_path):
        os.remove(video_path)
    
    if frames_count > 0:
        print(f"处理完成！共提取 {frames_count} 帧")
        print(f"帧图片保存在 {FRAMES_FOLDER} 目录下")
    else:
        print("处理失败")

if __name__ == '__main__':
    import sys
    if len(sys.argv) > 1 and sys.argv[1] == 'api':
        # 启动API服务
        app.run(host='0.0.0.0', port=5000, debug=True)
    else:
        # 运行命令行版本
        main()
